import {Colors, RedBlackNode} from "./RedBlackNode";
import AdelsonVelskiiLandiTree from "./AdelsonVelskiiLandiTree";
import {Compare, compareType, defaultCompare} from "../utils/Compare";

/**

 红黑树
 红黑树遵循以下规则
 1. 每个节点不是红的就是黑的
 2. 树的根节点是黑的
 3. 所有的叶子节点都是黑的 (用null表示引用的节点)
 4. 如果一个节点是红的, 那么它的两个子节点都是黑的;
 5. 不能有两个相邻的红节点, 一个红节点的父节点和子节点不能是红的;
 6. 从给定的节点到它的后代节点的所有路径包含同样数量的黑色节点

 ps: 如果一个节点的叔叔节点是黑色的, 那么 此时插入的节点会导致树失去平衡


 */

export default class RedBlackTree<T> extends AdelsonVelskiiLandiTree<T> {
  root: RedBlackNode<T> | null;

  constructor(compareFn: compareType<T> = defaultCompare) {
    super(compareFn);
    this.root = null;
    this.compareFn = compareFn;
  }

  insert(key: T) {
    if (this.root == null) {
      this.root = new RedBlackNode(key);
      this.root.color = Colors.Black;
    } else {
      const newNode = this.insertNode(this.root, key);
      this.fixTreeProperties(newNode);
    }
  }

  /**
   * 红黑树专用插入节点的操作
   * @param node
   * @param key
   */
  protected insertNode(node: RedBlackNode<T>, key: T): RedBlackNode<T> {
    if (this.compareFn(key, node.key) === Compare.NORMAL_THAN) {
      return node;
    }
    let index: "left" | "right" = this.compareFn(key, node.key) === Compare.LESS_THAN ? "left" : "right";

    if (node[index] === null) {
      node[index] = new RedBlackNode<T>(key);
      (<RedBlackNode<T>>node[index]).parent = node;
      return (<RedBlackNode<T>>node[index]);
    } else {
      return this.insertNode((<RedBlackNode<T>>node[index]), key);
    }

  }

  /**
   * 检查是否需要平衡
   * @param node
   */
  fixTreeProperties(node: RedBlackNode<T>) {
    while (node && node.parent && node.parent.isRed() && !node.isRed()) {
      let parentNode = node.parent;
      let grandParentNode = parentNode.parent;
      // 情况1 : 父节点是左侧节点
      if (grandParentNode && grandParentNode.left === parentNode) {
        const uncle = grandParentNode.right;
        // 情况1.1 : 叔叔节点也是红色 则只执行变色
        if (uncle && uncle.isRed()) {
          uncle.color = Colors.Black;
          parentNode.color = Colors.Black;
          grandParentNode.color = Colors.Red;
          node = grandParentNode;
        } else {
          //  叔叔节点是黑色, 则需要进行平衡
          // 父节点在左侧
          // 情况1.2: 节点是右侧子节点 -- 左旋转
          if (parentNode.right === node) {
            // 进行左旋转
            this.rotationRR(parentNode)
            node = parentNode;
            parentNode = <RedBlackNode<T>>node.parent;
          }
          // 情况1.3: 节点是左侧子节点 -- 右旋转
          this.rotationLL(grandParentNode)
          grandParentNode.color = Colors.Red;
          parentNode.color = Colors.Black;
          node = parentNode;
        }
      } else if (grandParentNode) {
        // 父节点为右侧节点
        const uncle = grandParentNode.left;
        if (uncle && uncle.isRed()) {
          uncle.color = Colors.Black;
          parentNode.color = Colors.Black;
          grandParentNode.color = Colors.Red;
          node = grandParentNode;
        } else {
          //  叔叔节点是黑色, 则需要进行平衡
          // 情况2.2: 节点是左侧子节点 -- 右旋转
          if (parentNode.left === node) {
            // 进行右旋转
            this.rotationLL(parentNode)
            node = parentNode;
            parentNode = <RedBlackNode<T>>node.parent;
          }
          // 情况2.3: 节点是左侧子节点 -- 右旋转
          this.rotationRR(grandParentNode)
          parentNode.color = Colors.Black;
          grandParentNode.color = Colors.Red;
          node = parentNode;
        }
      }
    }
    this.root && (this.root.color = Colors.Black)
  }


  rotationLL(node: RedBlackNode<T>): RedBlackNode<T> {
    const tmp = <RedBlackNode<T>>node.left;
    node.left = <null | RedBlackNode<T>>tmp.right;
    if (tmp?.right?.key){
      tmp.right.parent = node;
    }
    tmp.parent = node.parent;
    if(!node.parent){
      this.root = tmp;
    }
    else{
      if(node === node.parent.left){
        node.parent.left = tmp;
      }else{
        node.parent.right = tmp;
      }
    }
    tmp.right = node;
    node.parent = tmp;
    return tmp;
  }

  rotationRR(node: RedBlackNode<T>): RedBlackNode<T> {
     const tmp = <RedBlackNode<T>>node.right;
     node.right = <null | RedBlackNode<T>>tmp.left;
     if (tmp?.left?.key){
       tmp.left.parent = node;
     }
     tmp.parent = node.parent;
     if(!node.parent){
       this.root = tmp;
     }else{
       if(node.parent.left === node){
         node.parent.left = tmp;
       }else{
         node.parent.right = tmp;
       }
     }
     tmp.left = node;
     node.parent = tmp;
     return tmp;
  }
}
